########################################################################
# sftp.py - SFTP client based on Paramiko, for sending files to a
#	remote server.
#
# Created 2006.05.08 by Warren Young
#
# Copyright (c) 2006 by Pixel Corps.  This program may be used under
# the terms of the X11 license, a copy of which should have accompanied
# this program, in the LICENSE file.
########################################################################

import dialogs, fncheck, gui

import paramiko
import paramiko.hostkeys

import os, posixpath, thread

from stat import *


#### class SFTPUploader ################################################
# Defines the class that manages the SFTP upload process.

class SFTPUploader:
	"""
	Class to manage the SFTP upload process. It is the interface between
	the parts of the program that gather input from the user about what 
	to upload, and Paramiko which does the actual file transfer.
	"""

	chunkSize = 4096

	def __init__(self):
		"""
		Initialize the object, setting all member variables to their
		default values.
		"""

		self.filePickerGUI = False
		self.hostKeys = None
		self.hostName = None
		self.localFiles = []
		self.port = 22
		self.remoteDirRaw = None
		self.remoteDirVar = None
		self.remoteFiles = []
		self.remoteSubdirs = []
		self.sftp = None
		self.transport = None
		self.userName = None
		self.modOnly = False


	def __str__(self):
		"""
		@rtype: string
		@return: a concise string describing the object's current state.
		Does not tell all; useful mainly for debug strings.
		"""
	
		return 'sending files [%s] to sftp://%s@%s:%d/%s' % (
				', '.join(self.localFiles), self.userName,
				self.hostName, self.port, self.RemoteDir())


	def BuildRemotePaths(self):
		"""
		Go through file list looking for subdir qualifiers, building
		final set of remote paths.
		"""

		self.remoteFiles = []
		self.remoteSubDirs = []

		for f in self.localFiles:
			if f.find('::') >= 0:
				(d, f) = f.split('::')
				self.remoteSubdirs.append(d)
				self.remoteFiles.append(posixpath.join(
						self.RemoteDir(), d, os.path.basename(f)))
				self.localFiles[-1] = f
			else:
				self.remoteSubdirs.append(None)
				self.remoteFiles.append(posixpath.join(
						self.RemoteDir(), os.path.basename(f)))
			#print 'SFTP: "%s" -> "%s"' % (f, self.remoteFiles[-1])


	def CloseTransport(self):
		'Shuts down the SSH transport gracefully.'

		if self.transport and self.transport.is_active():
			print 'SFTP: shutting SSH transport down...',
			try:
				i = 0
				self.transport.close()
				while self.transport.is_active() and i < 50:
					print ".",
					time.sleep(0.1)
					i += 1
			except:
				pass
			print "Done!"

		self.transport = None
		self.sftp = None


	def Connect(self):
		'''Open an SSH connection to the remote server and create
		an SFTP channel atop it.'''

		# Establish TCP-level connection with server
		self.transport = paramiko.Transport((self.hostName, self.port))
		self.transport.start_client()

		# Check host key
		if not self.hostKeys:
			self.GetHostKeys()
		remoteKey = self.transport.get_remote_server_key()
		rfp = self.FingerprintStr(remoteKey)
		keys = self.hostKeys.lookup(self.hostName)
		if keys:
			if keys[remoteKey.get_name()] != remoteKey:
				if dialogs.YesNoQuestion("Remote server's host key ",
						"does not match cached key. Fingerprint:\n\n",
						rfp, "\n\nAccept it anyway?").answer:
					self.hostKeys.add(self.hostName,
							remoteKey.get_name(), remoteKey)
					self.hostKeys.save()
				else:
					print 'SFTP: user rejected non-matching host key'
					return False
		else:
			if dialogs.YesNoQuestion("I don't know any host keys for ",
					self.hostName, ". Accept new key? ",
					"Fingerprint:\n\n", rfp).answer:
				self.hostKeys.add(self.hostName,
						remoteKey.get_name(), remoteKey)
				self.hostKeys.save()
			else:
				return False

		# Try to load private key for authentication
		key = None
		try:
			key = paramiko.RSAKey.from_private_key_file(\
					os.path.expanduser('~/.ssh/id_rsa'))
		except:
			try:
				key = paramiko.RSAKey.from_private_key_file(\
						os.path.expanduser('~/ssh/id_rsa'))
			except:
				try:
					key = paramiko.DSAKey.from_private_key_file(\
							os.path.expanduser('~/.ssh/id_dsa'))
				except:
					try:
						key = paramiko.DSAKey.from_private_key_file(\
								os.path.expanduser('~/ssh/id_dsa'))
					except:
						pass
	
		# Authenticate with server
		try:
			# Try public key authentication first
			if key:
				print "SFTP: Authenticating with key: %s %s %s" % \
						(key.get_name(), str(key.get_bits()),
						self.FingerprintStr(key))
				self.transport.auth_publickey(self.userName, key)
			else:
				raise paramiko.AuthenticationException
		except Exception, err:
			# No key or key failed, so try password authentication
			print "SFTP: Public key auth error %s" % (err)
			print "SFTP: Authenticating with password"
			if self.GetPassword() == None:
				return False
			self.transport.auth_password(username = self.userName,
					password = self.password)

		# Open the SFTP channel atop SSH connection.
		self.sftp = paramiko.SFTPClient.from_transport(self.transport)
		return True


	def CreateDirectories(self):
		'''Create the remote directory hierarchy needed to store all
		files we'll be uploading.'''

		for dir in self.DirectoryList():
			parent, current = posixpath.split(dir)
			if not parent:		# relative path
				parent = '.'

			print "SFTP: Seeking '%s' in '%s'" % (current, parent)
			if current in self.sftp.listdir(parent):
				print "SFTP: found remote subdir '%s'" % dir
			else:
				print "SFTP: creating remote subdir '%s'" % dir
				try:
					self.sftp.mkdir(dir)
					self.sftp.chmod(dir, 0775)
				except IOError, err:
					print "SFTP: mkdir '%s' failed: %s" % (dir, err)
					dialogs.MessageBox(
							"Failed to create '%s':\n\n%s" % (dir, err))
					return False

		print "SFTP: Remote directory hierarchy created."
		return True


	def DirectoryList(self):
		"""
		Returns a list of directories we need to try and create on
		the remote host, in the order they need to be created.
		
		@rtype: array of string
		@return: list of directory names on the remote host
		"""

		# Break main remote directory name up into a series of directory
		# names, and build a list of directories, increasingly deep,
		# that traces the path to the final subdir.
		dirList = []

		remoteDir = self.RemoteDir()
		if remoteDir[0] == '/':
			last = '/'
		else:
			last = None

		parts = remoteDir.split('/')
		if parts[0] == '/' or not parts[0]:
			parts.pop(0)

		for p in parts:
			if last:
				last = posixpath.join(last, p)
			else:
				last = p

			dirList.append(last)

		# Continue building list of remote directory names, only adding
		# each unique subdirectory qualifier to the end of the last
		# directory name in the previous list each time.
		processed = []
		for s in self.remoteSubdirs:
			if s and not s in processed:
				dirList.append(posixpath.join(last, s))
				processed.append(s)

		print 'SFTP: directories to create:', str(dirList)
		return dirList


	def FileIsNew(self, local, remote):
		"""
		Checks if the local file is newer than the remote version

		@type local: string
		@param local: path to the local version of the file

		@type remote: string
		@param remote: path to the remote version of the file

		@rtype: boolean
		@return: True if the local copy is newer than the remote one
		"""

		# Get local file mod time
		lmod = os.stat(local)[ST_MTIME]

		if posixpath.split(remote)[1] in \
				self.sftp.listdir(posixpath.split(remote)[0]):
			rmod = self.sftp.stat(remote).st_mtime
		else:
			return True

		if lmod > rmod:
			return True

		return False	


	def FingerprintStr(self, key):
		"""
		Return the given key's fingerprint in a human-readable form
		
		@type key: paramiko.PKey
		@param key: a public Key in Paramiko's binary format

		@rtype: string
		@return: the key's fingerprint in colon-separated hex form
		"""
	
		bytes = []
		for c in key.get_fingerprint():
			bytes.append('%02X' % ord(c))
		return ':'.join(bytes)


	def GetHostKeys(self):
		"""
		Load the local "known hosts" file, containing the host keys of
		machines we've connected to successfully before.
		"""
	
		try:
			self.hostKeys = paramiko.hostkeys.HostKeys(\
					os.path.expanduser('~/.ssh/known_hosts'))
			self.sshDir = os.path.expanduser('~/.ssh')
		except IOError:
			try:
				self.hostKeys = paramiko.hostkeys.HostKeys(\
						os.path.expanduser('~/ssh/known_hosts'))
				self.sshDir = os.path.expanduser('~/ssh')
			except IOError:
				self.hostKeys = paramiko.hostkeys.HostKeys()
				if os.name == 'mac' or os.name == 'posix':
					self.sshDir = os.path.expanduser('~/.ssh')
				else:
					self.sshDir = os.path.expanduser('~/ssh')

				if not os.path.exists(self.sshDir):
					os.mkdir(self.sshDir, 0700)
		
		numFound = len(self.hostKeys.items())
		if numFound:
			print 'SFTP: loaded %d host keys' % numFound
		else:
			print 'SFTP: no host keys found in %s' % self.sshDir


	def GetPassword(self):
		"""
		Pop up a password request dialog to get the remote user's
		password.

		@rtype: string
		@return: the password, or None if the user cancelled the dialog
		"""

		self.password = dialogs.PasswordRequest('Password for %s@%s:%d' % (
				self.userName, self.hostName, self.port)).password
		return self.password


	def Missing(self):
		"""
		@rtype: string
		@return: a string containing a description of the first member
		variable that is unset and so will prevent us from trying the
		upload, or None if all are set.
		"""
		
		if not self.localFiles: return 'local file name'
		if not self.hostName: return 'host name'
		if not self.remoteDirRaw: return 'remote directory'
		if not self.userName: return 'user name'
		return None


	def RemoteDir(self):
		"""
		@rtype: string
		@return: the remote directory name, with remote directory
		variable substitution if needed
		"""
	
		dir = self.remoteDirRaw
		if dir and '%s' in dir:
			return dir % self.remoteDirVar
		else:
			return dir


	def StartUpload(self):
		"Initiates the upload sequence."
	
		# Check the file names for sanity
		files = self.remoteFiles
		for fn, err in zip(files, fncheck.CheckFileNames(files)):
			if err and not dialogs.YesNoQuestion(
					os.path.basename(fn), " fails file name check:\n\n",
					err, "\n\nUpload anyway?").answer:
				return

		# File names are okay, so try to start uploading
		yellOnExit = True
		for loop in range(3):
			try:
				# Establish SFTP connection to remote server
				if self.Connect():
					# Connection open, so do the transfer
					yellOnExit = False
					if self.CreateDirectories():
						self.UploadFiles()
					break
				else:
					print 'SFTP: Connection failed'
					yellOnExit = False
					break
			except EOFError:
				print 'SFTP: upload aborted!'
				self.Shutdown()
			except Exception, err:
				print 'SFTP: Connection failure:', err, ' (type',
				print err.__class__, ')'
				dialogs.MessageBox('Failed to connect to server:\n\n', err)
				self.Shutdown()

		if yellOnExit:
			dialogs.MessageBox("Too many failures authenticating. ",
					"Exiting...")


	def Shutdown(self):
		"""Shuts the SFTP process down, and terminates the thread
		containing it."""
	
		print 'SFTP THREAD SHUTTING DOWN'
		gui.Shutdown()
		self.CloseTransport()
		thread.exit()


	def ThreadEntry(self):
		'''Called by GUI once it starts up, telling us it's okay to
		start sending the file(s).'''

		missing = self.Missing()
		if not missing:
			print 'SFTP THREAD STARTED:', str(self)
			try:
				self.StartUpload()
			except Exception, err:
				print "Unhandled exception caught in SFTP:", err
		else:
			dialogs.MessageBox('Initialization error: %s not given' % \
					(missing))

		# Files uploaded, or failed trying, so shut everything down
		self.Shutdown()


	def TransferFile(self, lfile, rfile, fileNum, fileSize):
		"""
		Send the file in pieces, updating the GUI progress dialog
		as we go.

		@type lfile: string
		@param lfile: path to local file

		@type rfile: string
		@param rfile: path to remote file

		@type fileNum: integer
		@param fileNum: index of file in the set of files to upload

		@type fileSize: integer
		@param fileSize: count of bytes in the file
		"""

		bytes = 0
		while True:
			chunk = lfile.read(self.chunkSize)
			if chunk:
				rfile.write(chunk)
				bytes = bytes + len(chunk)
				gui.UpdateProgress(fileNum, bytes * 100 // fileSize)
			else:
				rfile.close()
				break


	def UploadFiles(self):
		'''Iterates over the file list, setting up repeated calls to
		L{TransferFile}.'''
		
		fileNum = 0
		for lpath, rpath in zip(self.localFiles, self.remoteFiles):
			if not self.modOnly or self.FileIsNew(lpath, rpath):
				# Open local file for reading, and remote for writing
				print "SFTP: sending '%s' -> '%s'" % (lpath, rpath)
				lfile = open(lpath, 'rb', 0)
				rfile = self.sftp.open(rpath, 'w')

				# Send the file
				self.TransferFile(lfile, rfile, fileNum,
						os.path.getsize(lpath))
				lfile.close()
				rfile.close()

				# Make remote file read-write by owner and group so
				# group members can change or remove it.
				self.sftp.chmod(rpath, 0664)
				print "SFTP: uploaded file %d/%d successfully!" % (
						fileNum + 1, len(self.localFiles))
			else:
				print "SFTP: %s is not new, skipping" % lpath

			fileNum += 1

